import subprocess
import time
import threading
from functools import wraps
import os

def gpu_monitor(interval=1, filename='gpu_usage.log'):
    # remove existing log file
    if os.path.exists(filename):
        os.remove(filename)

    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # Event to signal the monitoring thread to stop
            stop_event = threading.Event()

            # Create the monitoring thread  
            def log_gpu_usage():
                with open(filename, 'a') as f:
                    while not stop_event.is_set():
                        try:
                            # Run nvidia-smi command to get GPU usage
                            result = subprocess.run(
                                ['nvidia-smi', '--id=0', '--query-gpu=utilization.gpu,memory.used,memory.total', '--format=csv,noheader,nounits'],
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE,
                                text=True
                            )
                            if result.returncode == 0:
                                timestamp = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
                                f.write(f"{timestamp}, {result.stdout.strip()}\n")
                                f.flush()
                            else:
                                f.write(f"Error: {result.stderr.strip()}\n")
                                f.flush()
                        except Exception as e:
                            f.write(f"Exception: {str(e)}\n")
                            f.flush()
                        
                        if stop_event.wait(interval):
                            break

            # Create the monitoring thread 
            monitor_thread = threading.Thread(target=log_gpu_usage) 
            monitor_thread.start()

            try:
                t_start = time.time()
                result = func(*args, **kwargs)    # run raw function
                t_end = time.time()
                print(f"[INFO]Running model cost time : {t_end - t_start}s")
            finally:
                stop_event.set()                  # signal the monitoring thread to stop
                monitor_thread.join()
                
        return wrapper
    return decorator

    